package com.feidee.fd.sml.algorithm.component.feature

import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.feature.Binarizer
import org.apache.spark.sql.DataFrame

/**
  * @Author JunxinWang
  * @Date 2019/3/26 18:53
  * @Description
  * @Reviewer  YongChen
  */
case class BinarizeEncoderParam(
                                 override val input_pt: String,
                                 override val output_pt: String,
                                 override val hive_table: String,
                                 override val flow_time: String,
                                 override val inputCol: String,
                                 override val outputCol: String,
                                 override val preserveCols: String,
                                 override val modelPath: String,
                                 // 用于对连续特征进行二值化的阈值参数。大于阈值的功能将被二值化为1.0。等于或小于阈值的特性将被二值化为0.0。默认值：0.0
                                 threshold: Double = 0.0
                               ) extends FeatureParam {

  def this() = this(null, null, null, null, "input", "features", null, null, 0.0)

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("threshold" -> threshold)
    map
  }
}


class BinarizeEncoder extends AbstractFeatureEncoder[BinarizeEncoderParam] {

  override def setUp(param: BinarizeEncoderParam, data: DataFrame): Array[PipelineStage] = {
    val binarizer: Binarizer = new Binarizer()
      .setInputCol(param.inputCol)
      .setOutputCol(param.outputCol)
      .setThreshold(param.threshold)

    Array(binarizer)
  }

}

object BinarizeEncoder {

  def apply(paramStr: String): Unit = {
    new BinarizeEncoder()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    BinarizeEncoder(args(0))
  }
}
